//
//  Encrept.m
//  EHCommunicate
//
//  Created by 史博远 on 16/5/31.
//  Copyright © 2016年 Gjoy. All rights reserved.
//

#import "Encrept.h"

@implementation AES128

- (void)AES_SetKey :(Byte *)key
{
    [self AES_KeyExpansion:key :_wKey];
}

- (Byte *)AES_CipherBlock :(Byte *)input
{
    Byte state[4][4];
    int i,r,c;
    
    for(r=0; r<4; r++)
    {
        for(c=0; c<4 ;c++)
        {
            state[r][c] = input[c*4+r];
        }
    }
    
    [self AES_AddRoundKey:state :_wKey[0]];
    
    for(i=1; i<=10; i++)
    {
        [self AES_SubBytes:state];
        [self AES_ShiftRows:state];
        
        if(i!=10)
        {
            [self AES_MixColumns:state];
        }
        
        [self AES_AddRoundKey:state :_wKey[i]];
    }
    
    for(r=0; r<4; r++)
    {
        for(c=0; c<4 ;c++)
        {
            input[c*4+r] = state[r][c];
        }
    }
    
    return input;
}

- (Byte *)AES_InvCipherBlock :(Byte *)input;
{
    Byte state[4][4];
    int i,r,c;
    
    for(r=0; r<4; r++)
    {
        for(c=0; c<4 ;c++)
        {
            state[r][c] = input[c*4+r];
        }
    }
    
    [self AES_AddRoundKey:state :_wKey[10]];
    
    for(i=9; i>=0; i--)
    {
        [self AES_InvShiftRows:state];
        [self AES_InvSubBytes:state];
        [self AES_AddRoundKey:state :_wKey[i]];
        
        if(i)
        {
            [self AES_InvMixColumns:state];
        }
    }
    
    for(r=0; r<4; r++)
    {
        for(c=0; c<4 ;c++)
        {
            input[c*4+r] = state[r][c];
        }
    }
    
    return input;
}

- (NSData *)AES_Cipher :(Byte *)input :(unsigned int)length;
{
    Byte* in = input;
    int i;
    if(!length)
    {
        while(*(in+length++));
        in = input;
    }
    for(i=0; i<length; i += 16)
    {
        [self AES_CipherBlock:in + i];
    }
    return [NSData dataWithBytes:input length:length];
}

- (NSData *)AES_InvCipher :(Byte *)input :(unsigned int)length;
{
    Byte* in = input;
    int i;
    for(i=0; i<length; i += 16)
    {
        [self AES_InvCipherBlock:in + i];
    }
    return [NSData dataWithBytes:input length:length];
}

- (void)AES_KeyExpansion :(Byte *)key :(Byte[][4][4]) w
{
    int i,j,r,c;
    Byte rc[] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36};
    Byte t[4];
    Byte temp;
    for(r=0; r<4; r++)
    {
        for(c=0; c<4; c++)
        {
            w[0][r][c] = key[r+c*4];
        }
    }
    for(i=1; i<=10; i++)
    {
        for(j=0; j<4; j++)
        {
            for(r=0; r<4; r++)
            {
                t[r] = j ? w[i][r][j-1] : w[i-1][r][3];
            }
            if(j == 0)
            {
                temp=t[0];
                for(r=0; r<3; r++)
                {
                    t[r] = Sbox[t[(r+1)%4]];
                }
                t[3] = Sbox[temp];
                t[0] ^= rc[i-1];
            }
            for(r=0; r<4; r++)
            {
                w[i][r][j] = w[i-1][r][j] ^ t[r];
            }
        }
    }
}

- (Byte)AES_FFmul :(Byte)a :(Byte)b
{
    Byte bw[4];
    Byte res=0;
    int i;
    bw[0] = b;
    for(i=1; i<4; i++)
    {
        bw[i] = bw[i-1]<<1;
        if(bw[i-1]&0x80)
        {
            bw[i]^=0x1b;
        }
    }
    for(i=0; i<4; i++)
    {
        if((a>>i)&0x01)
        {
            res ^= bw[i];
        }
    }
    return res;
}

- (void)AES_SubBytes :(Byte[][4])state
{
    int r,c;
    for(r=0; r<4; r++)
    {
        for(c=0; c<4; c++)
        {
            state[r][c] = Sbox[state[r][c]];
        }
    }
}

- (void)AES_ShiftRows :(Byte[][4])state
{
    Byte t[4];
    int r,c;
    for(r=1; r<4; r++)
    {
        for(c=0; c<4; c++)
        {
            
            t[c] = state[r][(c+r)%4];
        }
        for(c=0; c<4; c++)
        {
            state[r][c] = t[c];
        }
    }
}

- (void)AES_MixColumns :(Byte[][4])state
{
    Byte t[4];
    int r,c;
    for(c=0; c< 4; c++)
    {
        for(r=0; r<4; r++)
        {
            t[r] = state[r][c];
        }
        for(r=0; r<4; r++)
        {
            state[r][c] = [self AES_FFmul:0x02 :t[r]]
            ^ [self AES_FFmul:0x03 :t[(r+1)%4]]
            ^ [self AES_FFmul:0x01 :t[(r+2)%4]]
            ^ [self AES_FFmul:0x01 :t[(r+3)%4]];
        }
    }
}

- (void)AES_AddRoundKey :(Byte[][4])state :(Byte[][4])k
{
    int r,c;
    for(c=0; c<4; c++)
    {
        for(r=0; r<4; r++)
        {
            state[r][c] ^= k[r][c];
        }
    }
}

- (void)AES_InvSubBytes :(Byte[][4])state
{
    int r,c;
    for(r=0; r<4; r++)
    {
        for(c=0; c<4; c++)
        {
            state[r][c] = InvSbox[state[r][c]];
        }
    }
}

- (void)AES_InvShiftRows :(Byte[][4])state
{
    Byte t[4];
    int r,c;
    for(r=1; r<4; r++)
    {
        for(c=0; c<4; c++)
        {
            t[c] = state[r][(c-r+4)%4];
        }
        for(c=0; c<4; c++)
        {
            state[r][c] = t[c];
        }
    }
}

- (void)AES_InvMixColumns :(Byte[][4])state
{
    Byte t[4];
    int r,c;
    for(c=0; c< 4; c++)
    {
        for(r=0; r<4; r++)
        {
            t[r] = state[r][c];
            
        }
        for(r=0; r<4; r++)
        {
            state[r][c] = [self AES_FFmul:0x0e :t[r]]
            ^ [self AES_FFmul:0x0b :t[(r+1)%4]]
            ^ [self AES_FFmul:0x0d :t[(r+2)%4]]
            ^ [self AES_FFmul:0x09 :t[(r+3)%4]];
        }
    }
}

@end
